#!/usr/bin/env python3
"""
一键运行 PeiYangGuide API 测试脚本
"""

import os
import sys
import subprocess
import argparse
import time
import signal

# 定义进程组
processes = []

def signal_handler(sig, frame):
    """处理 Ctrl+C 信号，确保所有进程都被终止"""
    print("\n正在停止所有进程...")
    for process in processes:
        if process.poll() is None:  # 如果进程还在运行
            process.terminate()
    sys.exit(0)

def main():
    parser = argparse.ArgumentParser(description="运行 PeiYangGuide API 测试")
    parser.add_argument("--no-server", action="store_true", help="不启动服务器，假设服务器已在运行")
    parser.add_argument("--no-data", action="store_true", help="不准备测试数据")
    parser.add_argument("-v", "--verbose", action="store_true", help="显示详细输出")
    parser.add_argument("-k", "--keyword", help="只运行包含特定关键字的测试")
    parser.add_argument("--test-db", action="store_true", help="使用SQLite测试数据库")
    parser.add_argument("--memory-db", action="store_true", help="使用SQLite内存数据库")
    parser.add_argument("--db-url", help="指定数据库URL")
    parser.add_argument("--failfast", action="store_true", help="首次测试失败后立即停止")
    args = parser.parse_args()

    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)

    # 设置环境变量
    env = os.environ.copy()
    if args.memory_db:
        env["DATABASE_URL"] = "sqlite:///:memory:"
        print("使用SQLite内存数据库进行测试")
    elif args.test_db:
        env["DATABASE_URL"] = "sqlite:///./test.db"
        print(f"使用SQLite文件数据库进行测试: {os.path.abspath('./test.db')}")
    elif args.db_url:
        env["DATABASE_URL"] = args.db_url
        print(f"使用自定义数据库连接: {args.db_url}")
    else:
        print("使用配置文件中定义的默认数据库连接")

    try:
        # 1. 准备测试数据（如果需要）
        if not args.no_data and not args.memory_db:  # 内存数据库不需要准备数据，因为每次都是全新的
            print("正在准备测试数据...")
            try:
                data_prep_process = subprocess.run(
                    [sys.executable, "tests/prepare_test_data.py"],
                    env=env,
                    check=False,
                    capture_output=not args.verbose
                )
                if data_prep_process.returncode == 0:
                    print("测试数据准备完成！")
                else:
                    print("测试数据准备失败，但将继续执行测试。")
                    if args.verbose and data_prep_process.stderr:
                        print("错误信息:")
                        print(data_prep_process.stderr.decode('utf-8'))
            except Exception as e:
                print(f"准备测试数据时出错: {e}")
                print("将继续执行测试，但某些测试可能会失败。")

        # 2. 启动API服务（如果需要）
        if not args.no_server:
            print("正在启动API服务...")
            try:
                server_cmd = [sys.executable, "run.py"]
                
                # 如果使用内存数据库，需要创建表结构
                if args.memory_db:
                    print("使用内存数据库时需要创建表结构...")
                    # 这里可以添加创建表结构的代码
                    env["CREATE_TABLES"] = "1"
                
                server_process = subprocess.Popen(
                    server_cmd,
                    env=env,
                    stdout=subprocess.PIPE if not args.verbose else None,
                    stderr=subprocess.PIPE if not args.verbose else None
                )
                processes.append(server_process)
                
                # 等待服务器启动
                print("等待服务器启动...")
                time.sleep(5)  # 增加等待时间
                
                # 检查服务器是否成功启动
                if server_process.poll() is not None:
                    print("警告：API服务启动失败，但将继续执行测试。")
                    if args.verbose and server_process.stderr:
                        print("错误信息:")
                        print(server_process.stderr.read().decode('utf-8'))
                else:
                    print("API服务已启动！")
            except Exception as e:
                print(f"启动API服务时出错: {e}")
                print("将继续执行测试，但某些测试可能会失败。")

        # 3. 运行测试
        print("开始运行测试...")
        pytest_cmd = [sys.executable, "-m", "pytest", "-v", "tests/api/"]
        if args.keyword:
            pytest_cmd.extend(["-k", args.keyword])
        if args.failfast:
            pytest_cmd.append("-xvs")
        
        test_process = subprocess.run(pytest_cmd, env=env)
        
        # 4. 处理测试结果
        if test_process.returncode == 0:
            print("所有测试通过！")
        else:
            print(f"测试执行失败，退出码: {test_process.returncode}")
        
        return test_process.returncode
    
    finally:
        # 5. 清理资源
        if not args.no_server:
            print("正在停止API服务...")
            for process in processes:
                if process.poll() is None:  # 如果进程还在运行
                    process.terminate()
                    process.wait(timeout=5)  # 等待进程正常结束
            print("API服务已停止！")

if __name__ == "__main__":
    sys.exit(main()) 